import py_trees
import random
import time


class checkBattery(py_trees.behaviour.Behaviour):
    def __init__(self, name):
        super().__init__(name)
        self.battery_level = 50
        self.battery_threshold = 30

    def update(self):
        self.battery_level -= random.randint(3, 5)
        if self.battery_level > self.battery_threshold:
            print(f"{self.battery_level} 电量充足，可以执行任务 ✓")
            return py_trees.common.Status.SUCCESS
        else:
            print(f"{self.battery_level} 电量不足，需要充电 ✗")
            return py_trees.common.Status.FAILURE


class doTask(py_trees.behaviour.Behaviour):
    def __init__(self, name):
        super().__init__(name)
        self.name = name

    def update(self):
        print(f"{self.name} 执行任务...")
        time.sleep(1)
        print(f"{self.name} 执行任务完成---")
        return py_trees.common.Status.SUCCESS


class ChargeBattery(py_trees.behaviour.Behaviour):
    """充电行为节点"""

    def __init__(self, name):
        super().__init__(name)

    def update(self):
        print("🔋 正在充电中...")
        # 模拟充电过程
        time.sleep(1)
        print("⚡ 充电完成!")
        return py_trees.common.Status.SUCCESS


def create_behaviour_tree():
    """
    创建基于电量条件的行为树
    结构：
    - 选择器（根节点）
      |- 序列：电量检查 + 执行任务
      |- 序列：电量不足 + 充电
    """
    root = py_trees.composites.Selector("电量决策", memory=False)
    # 重要: 创建共享的电量检查节点实例
    battery_check = checkBattery("电量检查")

    # 分支1：电量充足，执行任务
    task_branch = py_trees.composites.Sequence("执行任务分支", memory=True)
    # 行为节点
    task_branch.add_children(
        [
            battery_check,
            doTask("执行任务"),
        ]
    )

    # 分支2：电量不足，充电
    charge_branch = py_trees.composites.Sequence("充电分支", memory=True)
    charge_branch.add_children(
        [
            py_trees.decorators.Inverter("电量不足", battery_check),
            ChargeBattery("充电"),
        ]
    )

    root.add_children(
        [
            task_branch,
            charge_branch,
        ]
    )

    return root


def main():
    behaviour_tree = create_behaviour_tree()
    behaviour_tree.setup(timeout=15)
    # 打印树结构
    print("行为树结构:")
    print(py_trees.display.ascii_tree(behaviour_tree))

    count = 1
    while count < 10:
        print(f"第{count}次循环")
        behaviour_tree.tick_once()
        count += 1
        time.sleep(1)


if __name__ == "__main__":
    main()
